#!/usr/bin/env python
#encoding: utf-8
#Author: liwei
#Date: 2017-04-01
#File: psutil.py
#Env: python 2.6.6
# 手动实现并简化了psutil库 目的: 学习和练习

import os
import time
from collections import namedtuple

# 定义Linux下proc文件系统的目录
PROCFS_PATH = '/proc/'


###############################################
#                  公共接口                   #
###############################################

# 定义求使用占总量的百分比(参数：int/float)
def usage_percent(used, total, _round=None):
    try:
        ret = (float(used) / total) * 100    # 这个地方首先将used转成float, 否则整数相除为0
    except ZeroDivisionError:
        ret = 0.0 if isinstance(used, float) or isinstance(total, float) else 0
    if _round is not None:
        return round(ret, _round)
    else:
        return ret


###############################################
#                 Memory相关接口              #
###############################################

# 获取'/proc/meminfo'的字段信息, 返回字典类型格式
def _get_proc_meminfo_field():
    memory_field = {}
    with open(os.path.join(PROCFS_PATH, 'meminfo'), 'rb') as f:
        for line in f:
            fields = line.split()
            memory_field[fields[0]] = int(fields[1]) * 1024    # 将获取到的kb单位转换为B
    return memory_field


# 虚拟内存相关
def virtual_memory():
    # 获取meminfo相关信息存入字典
    memory_field = _get_proc_meminfo_field()
    
    # 取得/proc/meminfo各个字段的key-value到memory_field后
    total  = memory_field['MemTotal:']           # 总内存
    free   = memory_field['MemFree:']            # 剩余(不包括buffer/cache)
    buffer = memory_field['Buffers:']            # buffer
    cache  =  memory_field['Cached:']            # cache
    used   = total - free - buffer - cache       # 已经使用(不包括buffer/cache)
    if used < 0:
        used = total - free

    avail    = free + buffer + cache             # 剩余可用的(区别于free)
    shared   = memory_field['Shmem:']            # 共享内存
    active   = memory_field['Active:']           # 最近经常使用的内存
    inactive = memory_field['Inactive:']         # 最近不经常使用内存
    percent  = usage_percent(used, total, _round = 1) # 使用百分比 

    svmem = namedtuple('svmem', ['total', 'avail', 'percent', 'used', 'free', 
                       'active', 'inactive', 'buffer', 'cache', 'shared'])
    # 返回一个namedtuple
    return svmem(total, avail, percent, used, free, active, inactive, buffer, cache, shared)


# swap内存相关
def swap_memory():
    memory_field = _get_proc_meminfo_field()
    total = memory_field['SwapTotal:']
    free  = memory_field['SwapFree:']
    used  = total - free
    if used < 0:
        used = 0
    percent = usage_percent(used, total, _round = 1)
    
    # swap_in swap_out
    sin = sout = None
    ''' /proc/vmstat '''
    with open(os.path.join(PROCFS_PATH, 'vmstat'), 'rb') as f:
        for line in f:
            if line.startswith('pswpin'):
                sin = int(line.split()[1]) * 4 * 1024
            elif line.startswith('pswpout'):
                sout = int(line.split()[1]) * 4 * 1024
            if sin is not None and sout is not None:
                break
        # 如果存在sin或sout没有读取到的情况, 直接赋值为0
        else:
            sin  = 0
            sout = 0
    sswap = namedtuple('sswap', ['total', 'used', 'free', 'percent', 'sin', 'sout'])
    return sswap(total, used, free, percent, sin, sout)


###############################################
#                 CPU相关接口                 #
###############################################

# 计算总CPU各个字段的时间片值
def cpu_times():
    CLOCK_TICKS = os.sysconf("SC_CLK_TCK")    # 表示多少滴答/sec, 转换单位为秒s

    # values: cpu  6001656 3258 2677113 992372302 44464 10 5112 0 0
    with open(os.path.join(PROCFS_PATH, 'stat'), 'r') as f:                    # /proc/stat
        values = f.readline().split()[1:]
    fields = ['user', 'nice', 'system', 'idle', 'iowait', 'irq', 'softirq']    # 这7个字段是固定的

    # 不同的内核版本, 字段是有区别的.
    vlen = len(values)
    if vlen >= 8:    # Linux >= 2.6.11
        fields.append('steal')
    if vlen >= 9:    # Linux >= 2.6.24
        fields.append('guest')
    if vlen >= 10:   # Linux >= 3.2.0
        fields.append('guest_nice')

    # 定义namedtuple(scputimes)
    scputimes = namedtuple('scputimes', fields)
    values = [float(i) / CLOCK_TICKS for i in values]   # 时间片/CLOCK_TICKS
    return scputimes(*values)    # 返回namedtuple


# 获取具体每个cpu核心的各字段的时间片值
def per_cpu_times():
    CLOCK_TICKS = os.sysconf("SC_CLK_TCK")

    # /proc/stat 相关内容格式如下：
    #cpu  2405 0 4105 1129515 4564 0 361 0 0 0
    #cpu0 494 0 1236 282002 1035 0 24 0 0 0
    #cpu1 472 0 928 282650 1280 0 35 0 0 0
    #cpu2 420 0 947 282672 1036 0 143 0 0 0
    #cpu3 1017 0 993 282189 1212 0 158 0 0 0
    with open(os.path.join(PROCFS_PATH, 'stat'), 'r') as f:
        values = f.readline().split()[1:]
        fields = ['user', 'nice', 'system', 'idle', 'iowait', 'irq', 'softirq']
        
        # centos6 kernel=2.6.32  centos7 kernel=3.10.0
        vlen = len(values)
        if vlen >= 8:
            fields.append('steal')
        if vlen >= 9:
            fields.append('guest')
        if vlen >=10:
            fields.append('guest_nice')
    
        # define namedtuple(scputimes)
        scputimes = namedtuple('scputimes', fields)
    
        # 把各个cpu核心的数据存放到列表中.
        cpus = []
        for line in f:
            if line.startswith('cpu'):
                values = line.split()[1:]
                values = [float(i) / CLOCK_TICKS for i in values]
                cpus.append(scputimes(*values))
        return cpus


# 获取主机逻辑cpu的个数(process)
def cpu_count_logical():
    num = 0
    with open(os.path.join(PROCFS_PATH, 'cpuinfo'), 'r') as f:
        for line in f:
            if line.lower().startswith('process'):
                num += 1
    if num == 0:
        return None

    return num


# 获取主机物理cpu的个数(physical id)
def cpu_count_physical():
    mapping = {}
    current_info = {}

    # 计算各个物理CPU的核心数, 求和sum. 思路：以物理核id为key, 取对应的cpu cores为value.
    with open(os.path.join(PROCFS_PATH, 'cpuinfo'), 'r') as f:
        for line in f:
            line = line.strip().lower()
            if not line:    # 遇到空行(遇到完整的section后)
                if 'physical id' in current_info and 'cpu cores' in current_info:
                    mapping[current_info['physical id']] = current_info['cpu cores']
                    current_info = {}    # 处理完一个section后将current_info清零.
            else:           # 正在处理section
                if line.startswith('physical id') or line.startswith('cpu cores'):
                    key, value =  line.split('\t:', 1)    # split()的第二个参数1, 表示只分割1个段, 后边内容为一个段
                    current_info[key] = int(value)
    return sum(mapping.values()) or None


# 获取主机cpu的中断/上下文切换/系统调用相关信息
def cpu_stats():
    ctx_switches = None
    interrupts = None
    soft_interrupts = None
    with open(os.path.join(PROCFS_PATH, 'stat'), 'r') as f:
        for line in f:
            if line.startswith('ctxt'):
                ctx_switches = int(line.split()[1])
            elif line.startswith('intr'):
                interrupts = int(line.split()[1])
            elif line.startswith('softirq'):
                soft_interrupts = int(line.split()[1])
            if ctx_switches is not None and interrupts is not None and soft_interrupts is not None:
                break

    syscalls = 0
    scpustats = namedtuple('scpustats', ['ctx_switches', 'interrupts', 'soft_interrupts', 'syscalls'])
    return scpustats(ctx_switches, interrupts, soft_interrupts, syscalls)

# 计算cpu各个指标的有效时间和
def _cpu_total_times():
    times = cpu_times()
    total = sum(times)
    return total - times.guest

def _cpu_busy_times():
    times = cpu_times()
    busy = _cpu_total_times() - times.idle - times.iowait
    return busy

# 获取主机CPU的使用率(整体)
# 获取两次数据, 间隔时间0.1s, 计算0.1s内的使用率
def cpu_percent():
    total_t1 = _cpu_total_times()
    busy_t1  = _cpu_busy_times()
    time.sleep(0.01)
    total_t2 = _cpu_total_times()
    busy_t2  = _cpu_busy_times()
    return round((busy_t2 - busy_t1) / (total_t2 - total_t1) * 100, 1)


# 获取主机CPU的times相关的使用率
def cpu_times_percent():
    pass


# 具体到每个cpu的使用率
def per_cpu_percent():
    pass


# 测试主函数
if __name__ == '__main__':
    mem = virtual_memory()
    swap = swap_memory()
    #print mem.total, mem.avail, mem.percent, mem.used, mem.free, mem.active, mem.inactive, mem.buffer, mem.cache, mem.shared
    #print swap.total, swap.used, swap.free, swap.percent, swap.sin, swap.sout
    ret = cpu_times()
    #print ret
    ret = per_cpu_times()
    #print ret
    cpu = cpu_count_logical()
    #print cpu
    physical = cpu_count_physical()
    #print physical
    stats = cpu_stats()
    #print stats
    print _cpu_total_times()
    print _cpu_busy_times()
    print cpu_percent()

